import numpy as np
import struct
import matplotlib.pyplot as plt
import random
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from sklearn.datasets import load_digits  # 加载手写数字识别数据
from sklearn.cluster import KMeans
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix,classification_report

def get_random_init_centroids(ts,k,c_type):
    #initial:随机选k个质心
    centroids = np.empty((k,ts.shape[1]))
    if c_type == 1:
        ts_min, ts_max = ts.min(0), ts.max(0)
        #print("ts_min,ts_max",ts_min,ts_max)
        centroids = ts_min+(ts_max-ts_min)*np.random.rand(k,ts.shape[1])
    else:
        random_n = random.sample(range(0,ts.shape[0]),k)
        #print(random_n)
        for i in range(k):
            centroids[i] = ts[random_n[i]]
    return centroids

def assign_and_centroids(ts,last_centroids,k):
    print("===========================================================================assign_center===")
    print("last_centroids",last_centroids)

    distance = np.empty((ts.shape[0],k))
    assign = np.zeros((k,ts.shape[0]))
    min_dis = 9999
    min_pos = 0
    for i in range(ts.shape[0]):
        for j in range(k):
            a = np.linalg.norm(ts[i]-last_centroids[j])
            if j == 0 or a<min_dis:
                min_dis = a
                min_pos = j
            distance[i,j] = a
        #print("min_dis,min_pos",min_dis,min_pos)
        assign[min_pos,i] = int(min_pos+1)
    #print(distance)
    #print(assign)

    new_centroids = np.empty((k,ts.shape[1]))
    for i in range(k):
        b = np.zeros(ts.shape[1])
        #print("init b",b)
        count = 0
        for j in range(ts.shape[0]):
            if assign[i,j] != 0:
                b=b+ts[j]
                count+=1
                #print("b",b)
                #print("count",count)
        if count == 0:
            new_centroids[i]= last_centroids[i]
            print("no assignment")
        else:
            new_centroids[i]= b/count
    #print("new_centroids",new_centroids)
    return new_centroids,assign

def get_highest_SSE_cluster(ts,assign,new_centroids,k):
    max_sse = 0
    for i in range(k):
        sse=0
        for j in range(ts.shape[0]):
            if assign[i,j] != 0:
                #print("ts[j]",ts[j])
                #print("new_centroids[i]",new_centroids[i])
                sse=sse+np.sum(np.square(ts[j]-new_centroids[i]))
                #print("sse",sse)
        if sse>max_sse:
            max_sse = sse
            max_sse_cluster_label = i
    return max_sse_cluster_label

def get_highest_SSE_cluster_point(ts,assign,new_centroids,max_sse_cluster_label):
    max_sse = 0
    sse = 0
    for j in range(ts.shape[0]):
        if assign[max_sse_cluster_label,j] != 0:
            #print("ts[j]",ts[j])
            #print("new_centroids[i]",new_centroids[max_sse_cluster_label])
            sse=np.sum(np.square(ts[j]-new_centroids[max_sse_cluster_label]))
            #print("point_sse",sse)
        if sse>max_sse:
            max_sse = sse
            max_sse_cluster_label_point = ts[j]
    return max_sse_cluster_label_point

def repeat_2_getres(ts,last_centroids,k,max_iter):
    print("===========================================================================repeat===")
    t = 1
    while(t):
        #print("round",t,"========================")
        new_centroids,assign = assign_and_centroids(ts,last_centroids,k)
        #print("in round",t,"last_centroids=",last_centroids)
        #print("new_centroids=",new_centroids)
        #print("assign=",assign)
        if (new_centroids == last_centroids).all() or t>max_iter:
            print("same!")
            print("last_centroids",last_centroids)
            print("new_centroids ",new_centroids)
            break
        t = t+1
        last_centroids = new_centroids
    print("final t=",t)
    return new_centroids,assign

def empty_clusters(ts,final_centroids_unchecked,k,final_assign_unchecked):
    flag_empty = 0
    for i in range(k):
        print("np.max(final_assign_unchecked[i]),i",np.max(final_assign_unchecked[i]),i)
        if np.max(final_assign_unchecked[i]) == 0:
            print("there is a empty cluster")
            flag_empty = 1
            max_sse_cluster_label = get_highest_SSE_cluster(ts,final_assign_unchecked,final_centroids_unchecked,k)
            print("max_sse_cluster_label",max_sse_cluster_label)
            max_sse_cluster_label_point = get_highest_SSE_cluster_point(ts,final_assign_unchecked,final_centroids_unchecked,max_sse_cluster_label)
            print("max_sse_cluster_label_point",max_sse_cluster_label_point)
            final_centroids_unchecked[i] = max_sse_cluster_label_point
            final_centroids_checked = final_centroids_unchecked
    if flag_empty == 1:
        print("after sse new_centroids",final_centroids_checked)
        centroids,assign = repeat_2_getres(ts,final_centroids_checked,k,max_iter)
        return empty_clusters(ts,centroids,k,assign)
    else:
        return final_centroids_unchecked,final_assign_unchecked


def my_kmeans(ts,k,c_type,max_iter):
    
    #异常判断：类簇数要少于等于样本数
    if k>ts.shape[0]:
        #k = ts.shape[0]
        raise Exception("n_samples should >= n_clusters")
    #异常判断：c_type只能取1或2
    if type(c_type)!=int or (c_type!=1 and c_type!=2):
        #c_type = 1
        print(c_type)
        raise Exception("c_type should be int type of 1 or 2")
    
    #inital->初始化质心
    init_centroids = get_random_init_centroids(ts,k,c_type)
    print("init_centroids",init_centroids)

    #repeat->计算距离进行点分配+重新计算质心->until质心不再变化or到达最大迭代次数
    final_centroids_unchecked,final_assign_unchecked = repeat_2_getres(ts,init_centroids,k,max_iter)
    print("final centroid",final_centroids_unchecked)
    print("final assgin",final_assign_unchecked)
    
    #other-issue->处理空簇
    final_centroids,final_assign = empty_clusters(ts,final_centroids_unchecked,k,final_assign_unchecked)
        
    #结果规范化->转换成res_label形式
    res_label = np.zeros(ts.shape[0])
    for i in range(k):
        for d in range(final_assign.shape[1]):
            if final_assign[i,d]!=0:
                res_label[d] = final_assign[i,d]

    res_label = res_label.astype(np.int)
    print("res_label",res_label)
    return final_centroids,final_assign,res_label
    
if __name__ == '__main__':
    ts1 = np.array([[1,3,9,4,6],
                   [2,5,5,3,8],
                   [9,2,3,7,1]])
    ts2 = np.array([[0,0],
                   [1,0],
                   [0,1],
                   [1,1],
                   [3,3],
                   [4,3],
                   [3,4],
                   [4,4]])
    ts3 = np.array([[1,1],
                   [1,2],
                   [1,3],
                   [2,2],
                   [3,1],
                   [3,2],
                   [3,3]])
    
    digits = load_digits()
    ts = digits.data
    print("train_images1.shape",ts.shape)
    train_labels = digits.target
    print("train_labels.shape",train_labels.shape)
    print("Computing t-SNE embedding")
    tsne = TSNE(n_components=2, init='pca',random_state=0) #init='pca'
    train_image1_tsne = tsne.fit_transform(ts)
    print(train_image1_tsne)

    t_min, t_max = train_image1_tsne.min(0), train_image1_tsne.max(0)
    #print("t_min, t_max ",t_min, t_max)
    t_norm = (train_image1_tsne - t_min) / (t_max - t_min)  # 归一化
    print("t_norm",t_norm)
    
    k = 10
    max_iter = 200
    #print("ts,ts.shape",ts,ts.shape)
    c_type = 1#2
    #c_type = 1:随机生成k个样本点范围内inial质心
    #c_type = 2:随机从样本点中选择k个inial质心
    final_centroids,final_assign,res_label = my_kmeans(t_norm,k,c_type,max_iter)
    
    #print("right centroid",KMeans(n_clusters=k).fit_predict(ts))
    colors = ['FIREBRICK','lightcoral','coral','gold','springgreen',
              'darkgreen','skyblue','steelblue','darkgrey','darkorchid']
    plt.figure(figsize=(8, 8))
    for i in range(t_norm.shape[0]):
        plt.text(t_norm[i, 0], t_norm[i, 1], str(res_label[i]-1), color=colors[int(res_label[i])-1], 
                 fontdict={'weight': 'bold', 'size': 9})#plt.cm.Set1(int(label_pred[i])
    plt.scatter(final_centroids[:,0],final_centroids[:,1],c='k',marker='*')
    plt.xticks([])
    plt.yticks([])
    plt.show()

























